"""
展示高频预测，比较特性（x）与目标(y)的关系的可视化平台
"""
import datetime
import yaml
import matplotlib.pyplot as plt
import tkinter as tk
import pandas as pd
import numpy as np
from tkinter import ttk
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
from matplotlib.backend_bases import key_press_handler
from matplotlib.ticker import AutoLocator
from matplotlib.figure import Figure
from matplotlib import gridspec
from ini_data import get_data
from data_transformer import DataTransformer
from utils import cal_tztd

plt.rcParams['font.sans-serif'] = ['SimHei']  # 中文显示
plt.rcParams['axes.unicode_minus'] = False  # 负号显示


def load_config():
    with open("config.yaml", 'r', encoding='utf-8') as f:
        data = yaml.load(f.read(), yaml.FullLoader)
    print("loading_config ... >>> 读取本地配置")
    return data


def to_seasonal(df: pd.DataFrame, freq="月"):
    freq_map = {'日': "day", '周': 'week', '月': 'month', '季': 'quarter'}
    freq_ = freq_map[freq]
    df = pd.DataFrame(df)
    name = df.columns.to_list()[0]
    df.index = [pd.to_datetime(str(s)) for s in df.index]
    df['year'] = df.index.year
    if freq == "周":
        df['week'] = df.index.isocalendar().week
    elif freq == '月':
        df['month'] = df.index.month
    elif freq == '季':
        df['quarter'] = df.index.quarter
    else:
        df['day'] = df.index.dayofyear
    li = []
    for year, year_df in df.groupby(['year']):
        w_dic = {}
        for d_w_m, week_df in year_df.groupby([freq_]):
            w_dic[d_w_m] = week_df.mean(axis=0)[name]
        se = pd.Series(w_dic, name=year)
        li.append(se)
    data = pd.concat(li, axis=1)
    data.sort_index(inplace=True)
    return data


def ini_tree_x(tree: ttk.Treeview, x: dict, cfg: dict) -> dict:
    yoy_range = cfg['yoy_range']
    value_range = cfg['value_range']
    tree_dict = {}
    for first_i, (k, v) in enumerate(x.items()):
        tree_dict[k] = tree.insert("", first_i, k,
                                   # text=k,
                                   values=(k, v['频率']))
        for second_i, col in enumerate(['原始值',
                                        '当月值_平均差分',
                                        '当月值_最后',
                                        '当月值_平均',
                                        '累计转当月_同比',
                                        '当月值_季调',
                                        '当月环比_最后',
                                        '当月环比_平均',
                                        '当月环比_季调',
                                        '当季值_最后',
                                        '当季值_平均',
                                        '当月同比',
                                        '当季同比']):
            tree_dict[f"{k}>>{col}"] = tree.insert(tree_dict[k], second_i, f"{k}>>{col}",
                                                   text=col,
                                                   values=(col, '-'))
        for i in range(1, yoy_range):
            col = f"当月同比lag_{i}"
            second_i_ = i + second_i
            tree_dict[f"{k}>>{col}"] = tree.insert(tree_dict[k], second_i_, f"{k}>>{col}",
                                                   text=col,
                                                   values=(col, '-'))
        for i in range(1, value_range):
            col = f"当月值lag_{i}"
            second_ii_ = i + second_i_
            tree_dict[f"{k}>>{col}"] = tree.insert(tree_dict[k], second_ii_, f"{k}>>{col}",
                                                   text=col,
                                                   values=(col, '-'))
        for i in range(1, 6):
            col = f"当月同比ma_{i}"
            second_iii_ = i + second_ii_
            tree_dict[f"{k}>>{col}"] = tree.insert(tree_dict[k], second_iii_, f"{k}>>{col}",
                                                   text=col,
                                                   values=(col, '-'))
        for i in range(1, 14):
            col = f"当月值ma_{i}"
            second_iiii_ = i + second_iii_
            tree_dict[f"{k}>>{col}"] = tree.insert(tree_dict[k], second_iiii_, f"{k}>>{col}",
                                                   text=col,
                                                   values=(col, '-'))
    return tree_dict


class Gui:
    def __init__(self):
        super(Gui, self).__init__()
        self.cfg = load_config()
        self.root = tk.Tk()
        self.root.title(self.cfg['welcome_title'])
        self.root.geometry("1900x1100")
        self.bg = "#F07C82"  # 设置背景色
        self.s_f = ('Microsoft YaHei', 15,)  # second_font, 二级字体
        self.logo = plt.imread('meta/logo_.png')
        self.book = ttk.Notebook(self.root, height=950, width=1400)
        self.book.place(x=440, y=0)

        self.fm_args = {"relief": "ridge", "bd": 3, "bg": self.bg}
        self.fm_show = tk.Frame(self.book, height=1030, width=1400, **self.fm_args)
        self.fm_show.place(x=440, y=100)
        self.fm_seasonal = tk.Frame(self.book, height=630, width=1400, **self.fm_args)
        self.fm_seasonal.place(x=440, y=100)
        self.fm_scatter = tk.Frame(self.book, height=630, width=1400, **self.fm_args)
        self.fm_scatter.place(x=440, y=100)
        self.book.add(self.fm_show, text=self.cfg['sheet1_name'])
        self.book.add(self.fm_seasonal, text=self.cfg['sheet2_name'])
        self.book.add(self.fm_scatter, text=self.cfg['sheet3_name'])

        self.xy_dict = get_data(self.cfg['update'], self.cfg['path'])
        self.x_dic = {k: v for k, v in self.xy_dict.items() if v['x']}
        self.x_dic = {k: v for k, v in self.xy_dict.items() if v['频率'] in self.cfg['x_freqs']}
        self.y_dic = {k: v for k, v in self.xy_dict.items() if v['y']}

        self.tree_y = ttk.Treeview(self.root, columns=['Y', '频率'], height=20)
        self.tree_y.column("#0", width=10, anchor="center")
        self.tree_y.column('Y', width=300, minwidth=300)
        self.tree_y.heading('Y', text='Y')
        self.tree_y.column('频率', width=100, minwidth=100)
        self.tree_y.place(x=5, y=10)
        self.tree_y_dict = ini_tree_x(self.tree_y, self.y_dic, self.cfg)
        self.tree_y.bind("<<TreeviewSelect>>", self._tree_fun)

        self.tree_x = ttk.Treeview(self.root, columns=['X', '频率'], height=24)
        self.tree_x.column("#0", width=10, anchor="center")
        self.tree_x.column('X', width=300)
        self.tree_x.heading('X', text='X')
        self.tree_x.column('频率', width=100)
        self.tree_x.place(x=5, y=460)
        self.tree_x_dict = ini_tree_x(self.tree_x, self.x_dic, self.cfg)
        self.tree_x.bind("<<TreeviewSelect>>", self._tree_fun)

        self.y_name = list(self.y_dic)[0]
        self.x_name = list(self.x_dic)[0]

        self._ini_fig()

    def _on_key_press(self, event):
        key_press_handler(event, self.canvas, self.toolbar)

    def _on_key_press_seasonal(self, event):
        key_press_handler(event, self.canvas_seasonal, self.toolbar_seasonal)

    @staticmethod
    def _basic_ini_fig(fig_size, fm, key_press_func, col_row):
        fig = Figure(figsize=fig_size, dpi=78)
        canvas = FigureCanvasTkAgg(fig, fm)
        canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand=1)
        toolbar = NavigationToolbar2Tk(canvas, fm)
        toolbar.update()
        canvas.mpl_connect("key_press_event", key_press_func)
        spec = gridspec.GridSpec(ncols=col_row[0], nrows=col_row[1], figure=fig, wspace=0.85, left=0.08, right=0.92)
        return fig, canvas, toolbar, spec

    def _ini_fig(self):
        self.fig, self.canvas, self.toolbar, self.spec = self._basic_ini_fig(
            fig_size=(18, 11), fm=self.fm_show, key_press_func=self._on_key_press, col_row=(20, 20))
        self.fig_seasonal, self.canvas_seasonal, self.toolbar_seasonal, self.spec_seasonal = self._basic_ini_fig(
            fig_size=(18, 5), fm=self.fm_seasonal, key_press_func=self._on_key_press_seasonal, col_row=(20, 20))
        self.fig_scatter, self.canvas_scatter, self.toolbar_scatter, self.spec_scatter = self._basic_ini_fig(
            fig_size=(18, 11), fm=self.fm_scatter, key_press_func=self._on_key_press_seasonal, col_row=(20, 20))

    def _set_x_transformer(self, data_type, name, xy='x'):
        print('data_type', data_type)
        if xy == 'x':
            self.xy_transformer_data = DataTransformer(self.x_dic[name])
            self.x_data_type = data_type
            if '当月' in data_type:
                self.x_plot_freq = '月'
            else:
                self.x_plot_freq = self.x_dic[name]['频率']
        else:
            self.xy_transformer_data = DataTransformer(self.y_dic[name])
            self.y_data_type = data_type
            if '当月' in data_type:
                self.y_plot_freq = '月'
            else:
                self.y_plot_freq = self.y_dic[name]['频率']

        if data_type == '原始值':
            se = self.xy_transformer_data.data_dic['df']
        elif data_type == '当月值_平均差分':
            se = self.xy_transformer_data.data_dic['df'].diff(1)
        elif data_type == '当月值_最后':
            se = self.xy_transformer_data.to_freq('月', how='last')
        elif data_type == '当月值_平均':
            se = self.xy_transformer_data.to_freq('月', how='mean')
        elif data_type == '当月值_季调':
            se = self.xy_transformer_data.to_freq('月', how='mean')
            se = self.sea_adj(se)
        elif data_type == '累计转当月_同比':
            se = self.xy_transformer_data.cum_to_mon()
            print(se.shape)
        elif data_type == '当月环比_最后':
            se = self.xy_transformer_data.to_freq('月', how='last')
            se = se.pct_change(1)
        elif data_type == '当月环比_平均':
            se = self.xy_transformer_data.to_freq('月', how='mean')
            se = se.pct_change(1)
        elif data_type == '当月环比_季调':
            _, se = self.xy_transformer_data['月_平均_环比_lag_0']
            se = self.sea_adj(se)
        elif data_type == '当季值_最后':
            se = self.xy_transformer_data.to_freq('季', how='last')
        elif data_type == '当季值_平均':
            se = self.xy_transformer_data.to_freq('季', how='mean')
        elif data_type == '当月同比':
            se = self.xy_transformer_data.to_yoy('月')
        elif data_type == '当季同比':
            se = self.xy_transformer_data.to_yoy('季')
        elif 'lag' in data_type:
            v = int(data_type.split('lag_')[1])
            if '当月同比' in data_type:
                se = self.xy_transformer_data.to_yoy('月').shift(v, 'M')
            elif '当月值' in data_type:
                se = self.xy_transformer_data.to_freq('月').shift(v, 'M')
        elif 'ma' in data_type:
            v = int(data_type.split('ma_')[1])
            if '当月同比' in data_type:
                se = self.xy_transformer_data.to_yoy('月').rolling(v).mean()
                se = pd.DataFrame(se)
                se = se.fillna(method='bfill')
            elif '当月值' in data_type:
                se = self.xy_transformer_data.to_freq('月', how='last').rolling(v).mean()
                se = pd.DataFrame(se)
                se = se.fillna(method='bfill')

        return se.copy()

    def _tree_fun(self, event):
        print(event.widget.selection()[0])
        x_name_tp = self.tree_x.selection()
        y_name_tp = self.tree_y.selection()

        if (len(y_name_tp) > 0) and (len(x_name_tp) > 0):
            sp_x = x_name_tp[0].split(">>")
            sp_y = y_name_tp[0].split('>>')
            x_plot = False
            y_plot = False

            if len(sp_x) == 2 and len(sp_y) == 2:
                self.x_name, target_type_x = sp_x
                self.x_se = self._set_x_transformer(data_type=target_type_x, name=self.x_name, xy='x')
                x_plot = True

            if len(sp_y) == 2:
                self.y_name, target_type_y = sp_y
                self.y_se = self._set_x_transformer(data_type=target_type_y, name=self.y_name, xy='y')
                y_plot = True

            self._fig_plot(x_plot, y_plot)

    def _fig_plot(self, x_plot, y_plot):
        self.fig.clear()
        self.fig_seasonal.clear()
        self.fig_scatter.clear()
        self._basic_plot(x_plot, y_plot)
        self._seasonal_plot(x_plot, y_plot)
        self.canvas.draw()
        self.canvas_seasonal.draw()
        self.canvas_scatter.draw()

    def _basic_plot(self, x_plot, y_plot):
        if x_plot:
            ax_x = self._x_fig_plot()
        if x_plot and y_plot:
            ax_y = ax_x.twinx()
            self._y_fig_plot(ax_y, ax_x)
            # if ('月' in self.y_data_type) and ('月' in self.x_data_type):
            try:
                self._xy_scatter_plot()
            except:
                print('plot scatter failed')
        if self.cfg['download_y']:
            self.y_se.to_excel('y.xlsx')

    def _seasonal_plot(self, x_plot, y_plot):
        ax_logo = self.fig_seasonal.add_subplot(self.spec_seasonal[:2, :6])
        self.logo_plot(ax_logo)
        if y_plot:
            ax_y = self.fig_seasonal.add_subplot(self.spec_seasonal[3:15, :9])
            self._seasonal_fig_plot(ax_y, self.y_se, self.y_data_type, xy='y')
        if x_plot:
            ax_x = self.fig_seasonal.add_subplot(self.spec_seasonal[3:15, 11:])
            self._seasonal_fig_plot(ax_x, self.x_se, self.x_data_type, xy='x')

    def _x_fig_plot(self):
        x_d = self.x_dic[self.x_name]
        ax_text_x = self.fig.add_subplot(self.spec[3:4, :4])
        t_str = f"名称：{x_d['name']} 原始频率：{x_d['频率']}"
        ax_text_x.set_title(t_str)
        self.set_transparent(ax_text_x)
        ax_logo = self.fig.add_subplot(self.spec[6:9, :7])
        self.logo_plot(ax_logo)
        ax_x = self.fig.add_subplot(self.spec[9:, :])
        ax_x.plot(self.x_se, label=f'x >> {self.x_name}>>左轴')
        ax_x.legend(loc='lower left')
        return ax_x

    def _y_fig_plot(self, ax_y, ax_x):
        ax_y.plot(self.y_se, color='r', label=f'y >> {self.y_name}>>右轴')
        acc, acc_dic = cal_tztd(self.x_se.iloc[:, 0], self.y_se.iloc[:, 0], t_cut=self.cfg['t_cut'])
        t_str = f"{self.cfg['t_cut']}至今"

        ax_x.set_title(f"【{self.y_name}】 vs 【{self.x_name}】 同涨同跌率【{acc:4.2%}】 {t_str}")
        y_upper = np.max(self.y_se.iloc[:, 0])
        y_lower = np.min(self.y_se.iloc[:, 0])
        for t, (up_down, t_li) in acc_dic.items():
            t_start, t_end = t_li
            if up_down == 1:
                c = 'red'
            elif up_down == -1:
                c = 'green'
            else:
                c = 'gray'
            ax_y.fill_between([t_start, t_end], y1=y_upper, y2=y_lower, color=c, alpha=0.1)
        ax_y.legend(loc='best')

    def _xy_scatter_plot(self):
        ax_scatter = self.fig_scatter.add_subplot(self.spec[:, :])
        se_xy = pd.concat([self.x_se, self.y_se], axis=1)
        se_xy.dropna(axis=0, inplace=True)
        x_s = se_xy.iloc[:, 0].to_list()
        y_s = se_xy.iloc[:, 1].to_list()
        ax_scatter.scatter(x=x_s, y=y_s)
        ax_scatter.set_xlabel(xlabel=self.x_name, fontsize=20)
        ax_scatter.set_ylabel(ylabel=self.y_name, fontsize=20)

    def _seasonal_fig_plot(self, ax, se, data_type, xy):
        xy_map = {'x': self.x_plot_freq, 'y': self.y_plot_freq}
        if isinstance(se, pd.Series):
            name = se.name
        else:
            name = se.columns[0]
        df = pd.DataFrame(se.copy())
        df_seasonal = to_seasonal(df, xy_map[xy])
        today = datetime.date.today()
        for year, year_se in df_seasonal.iteritems():
            if year in self.cfg['color_map']:
                year_se.index = [f"{v}{xy_map[xy]}" for v in year_se.index]
                c, m = self.cfg['color_map'][year]
                ax.plot(year_se, label=f"{year}", marker=m, color=c)
                if year == 2022:
                    try:
                        for month in self.cfg['seasonal_marker']:
                            value = year_se[month]
                            if ('同比' in data_type) or ('环比' in data_type):
                                txt = "{:.1%}".format(value)
                            else:
                                txt = "{:.2f}".format(value)
                            ax.text(month, value, txt, fontsize=16, color='gray')
                    except:
                        pass
                ax.legend()
        ax.xaxis.set_major_locator(AutoLocator())
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.set_title(f"原始数据：{name} >> 数据类型：{data_type}")
        ax.yaxis.grid()

    def logo_plot(self, ax):
        ax.imshow(self.logo, )
        self.set_transparent(ax)

    @staticmethod
    def set_transparent(ax):
        for sp in ['right', 'top', 'left', 'bottom']:
            ax.spines[sp].set_visible(False)
        ax.xaxis.set_visible(False)
        ax.yaxis.set_visible(False)

    @staticmethod
    def sea_adj(df):
        assert len(df) > 60
        dic = {}
        for i in range(1, 13):
            dic[i] = df[df.index.month == i].mean().to_list()[0]
        m = np.mean(pd.Series(dic))
        factor = {k: v / m for k, v in dic.items()}
        name = df.columns[0]
        for idx in df.index:
            m_factor = factor[idx.month]
            df.loc[idx, name] = df.loc[idx, name] / m_factor
        return df


if __name__ == "__main__":
    from utils import list_excel

    list_excel()
    G = Gui()
    G.root.mainloop()
